package com.ruoyi.teaching.flink;

import com.ruoyi.common.core.utils.SpringUtils;
import com.ruoyi.teaching.domain.YaHighInterview;
import com.ruoyi.teaching.index.Interview;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.springframework.data.elasticsearch.core.ElasticsearchRestTemplate;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;

import static com.ruoyi.teaching.flink.MySQLSource.BATCH_SIZE;

/**
 * Sink 是数据流处理管道中的一个重要组件，负责将处理后的数据输出到外部系统或存储。Sink 是 Flink 流处理应用的终点，它定义了数据如何被写入到目标存储或系统中。
 */
@Slf4j
public class BatchESSink extends RichSinkFunction<YaHighInterview> {

    private int totalCount = 0;

    private transient ElasticsearchRestTemplate elasticsearchRestTemplate;

    private transient List<Interview> buffer;

    @Override
    public void open(Configuration parameters) {
        // 初始化缓冲区和Elasticsearch模板
        buffer = new ArrayList<>(BATCH_SIZE);
        elasticsearchRestTemplate = SpringUtils.getBean(ElasticsearchRestTemplate.class);
    }

    @Override
    public void invoke(YaHighInterview record, Context context) {
        // 将YaHighInterview对象转换为Interview对象
        Interview interview = new Interview();
        interview.setId(record.getId());
        interview.setQuestion(record.getQuestion());
        interview.setAnswer(record.getAnswer());
        interview.setKnowledgePoint(record.getKnowledgePoint());
        interview.setImportance(record.getImportance());
        interview.setStage(record.getStage());
        buffer.add(interview);
        // 当缓冲区达到批量大小时，写入Elasticsearch
        if (buffer.size() >= BATCH_SIZE) {
            writeBufferToES();
        }
    }

    private void writeBufferToES() {
        totalCount+=buffer.size();
        // 将缓冲区中的数据写入Elasticsearch
        log.debug("写入新批次数据到ES，批次大小：{}", buffer.size());
        elasticsearchRestTemplate.save(buffer);
        log.debug("已写入总记录数：{}", totalCount);
        buffer.clear();
    }

    @Override
    public void close() throws Exception {
        // 在关闭前，确保缓冲区中的剩余数据也被写入Elasticsearch
        if (!buffer.isEmpty()) {
            writeBufferToES();
        }
        log.debug("全量同步到ES结束");
    }
}
